import os
import numpy as np
import cv2 as cv
import torch
from seg_system.application_config import ApplicationConfig
from seg_system.segmentation.service.ceil.networks.csnet import CSNet as CeilCSNet
from seg_system.segmentation.service.csnet.CSNetService import CSNetService
from seg_common.annotation.CommonAnnotation import CommonAnnotation


class CSNetCeilService(CSNetService):
    """CSNetCeilService:
        阅读了源代码的部分，并进行了实验，新的细胞分割在旧CSNet上增加了一组功能性输出
        所以只要对输出部分进行改进即可

        需要增加细胞分割结果的保存路径
    """

    @CommonAnnotation.override()
    def after(self, obj, **kwargs):
        o_c, o_n = obj
        c_l, n_l = [], []

        for each_c, each_n in zip(o_c, o_n):
            each_c = self.return_prediction(each_c)
            each_n = self.return_prediction(each_n)

            c_l.append(each_c)
            n_l.append(each_n)

        return c_l, n_l

    @CommonAnnotation.override()
    def save(self, obj, path, name, **kwargs):
        c_o, n_o = obj
        c_o, n_o = self.to_list(c_o), self.to_list(n_o)
        name = self.to_list(name)

        each_from = kwargs.get('each_from', None)
        ceil_path = kwargs.get('ceil_path', None)
        tmp_path = kwargs.get('tmp_path', None)
        assert each_from is not None
        assert ceil_path is not None
        assert tmp_path is not None

        alpha = kwargs.get('alpha', 0.5)
        beta = 1.0 - alpha
        color = kwargs.get('color', np.array([0, 212, 255]))  # yellow bgr
        color_ceil = kwargs.get('color', np.array([255, 191, 0]))  # blue bgr

        for each_c_o, each_n_o, each_n in zip(c_o, n_o, name):
            file = os.path.join(path, each_n)
            ceil_path = os.path.join(ceil_path, each_n)
            tmp_file = os.path.join(tmp_path, each_n)

            ori_file = cv.imread(os.path.join(each_from, each_n))

            # 这里该不该看老师，现在是整合到原图上，或者就直接分开展示
            # 添加神经分割到原图
            no_zero_index = np.nonzero(each_n_o)
            ori_file[no_zero_index[0], no_zero_index[1], :] = \
                ori_file[no_zero_index[0], no_zero_index[1], :] * beta + alpha * color
            # 添加细胞分割到原图
            no_zero_index = np.nonzero(each_c_o)
            ori_file[no_zero_index[0], no_zero_index[1], :] = \
                ori_file[no_zero_index[0], no_zero_index[1], :] * beta + alpha * color_ceil

            # n代表神经，c代表细胞
            cv.imwrite(file, each_n_o)
            cv.imwrite(tmp_file, ori_file)
            # cv.imwrite(ceil_path, each_c_o)

            assert os.path.exists(file)
            assert os.path.exists(tmp_path)

    def __init__(self, model_path: str = ""):
        super().__init__(model_path)
        self.model = CeilCSNet(1, 1)
        model_path = os.path.join(model_path, ApplicationConfig.NetConfig.SEGMENTATION_DICT[
            ApplicationConfig.NetConfig.SEGMENTATION_USE_CSNET_CEIL])

        if model_path != "":
            model_d = torch.load(model_path)
            self.model.load_state_dict(model_d)
        self.model.to(ApplicationConfig.SystemConfig.DEVICE)
        self.model.eval()